from encoder import *
from torch.utils.data import Dataset, DataLoader
from einops import rearrange, reduce, repeat
import torch.nn as nn
import torch
import torch.nn.functional as F
import pytorch_lightning as pl
from pytorch_lightning.metrics import Accuracy
import numpy as np

class MyDataSet(Dataset):
    def __init__(self, x, y):
        super().__init__()
        self.x = x
        self.y = y
        self.len = x.shape[0]
        
    def __getitem__(self, idx):
        return self.x[idx, :], self.y[idx]
    
    def __len__(self):
        return self.len



def mish(input):

    return input * torch.tanh(F.softplus(input))

class Mish(nn.Module):

    def __init__(self):
        '''
        Init method.
        '''
        super().__init__()

    def forward(self, input):
        '''
        Forward pass of the function.
        '''
        return mish(input)

class MLPLayer(nn.Module):
    def __init__(self, dim_in, dim_out, res_coef = 0, dropout_p = 0.1):
        super().__init__()
        self.linear  = nn.Linear(dim_in, dim_out)
        self.res_coef = res_coef
        self.activation = Mish()
        self.dropout = nn.Dropout(dropout_p)
        self.ln = nn.LayerNorm(dim_out)
    
    def forward(self, x):
        y = self.linear(x)
        y = self.activation(y)
        y = self.dropout(y)
        if self.res_coef == 0:
            return self.ln(y)
        else:
            return self.ln(self.res_coef*x +y )

       
class MyNetwork(nn.Module):
    def __init__(self, dim_in, dim, res_coef=0.5, dropout_p = 0.1, n_layers = 10):
        super().__init__()
        self.mlp = nn.ModuleList()
        self.first_linear = MLPLayer(dim_in, dim)
        self.n_layers = n_layers
        for i in range(n_layers):
            self.mlp.append(MLPLayer(dim, dim, res_coef, dropout_p))
        self.final = nn.Linear(dim, 1)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x):
        x = self.first_linear(x)
        for layer in self.mlp:
            x = layer(x)
        x= self.sigmoid(self.final(x))
        return x.squeeze()


class TrainingModule(pl.LightningModule):
    def __init__(self, dim_in, dim, res_coef=0, dropout_p=0, n_layers=10):
        super().__init__()
        self.backbone = MyNetwork(dim_in, dim, res_coef, dropout_p, n_layers)
        self.loss = nn.BCELoss()
        self.accuracy = Accuracy()
    def forward(self, x):
        return self.backbone(x)
    def validation_step(self, batch, batch_idx):
        x, y = batch
        x = self.backbone(x)
        loss = self.loss(x, y.type(torch.float32))
        acc = self.accuracy(x, y)
        self.log("Validation loss", loss)
        self.log("Validation acc", acc)
        return loss, acc
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        x = self.backbone(x)
        loss = self.loss(x, y.type(torch.float32))
        acc = self.accuracy(x, y)
        self.log("Training loss", loss)
        self.log("Training acc", acc)
        
        return loss
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

import os
class CheckpointEveryNSteps(pl.Callback):
    def __init__(self, save_step_frequency):
        self.save_step_frequency = save_step_frequency

    def on_batch_end(self, trainer: pl.Trainer, _):
        epoch = trainer.current_epoch
        global_step = trainer.global_step
        if global_step % self.save_step_frequency == 0:
            filename = "epoch=" + str(epoch) + "_step=" + str(global_step)+".ckpt"
            ckpt_path = os.path.join(trainer.checkpoint_callback.dirpath, filename)
            trainer.save_checkpoint(ckpt_path)





class TabDataset(Dataset):
    def __init__(self, x, y):
        super().__init__()
        self.x = torch.from_numpy(x).type(torch.int32) 
        self.y = torch.from_numpy(y).type(torch.int32)

    def __getitem__(self, idx):
        return self.x[idx, :], self.y[idx]

    def __len__(self):
        return self.x.shape[0]

class EmbeddingFactory(nn.Module):
    def __init__(self, x, dim_out):
        super().__init__()
        self.dim_out = dim_out
        self.module_list = nn.ModuleList(
            [nn.Embedding(len(set(np.unique(x[:, col]))), dim_out) for col in range(x.shape[1])])

    def forward(self, x):
        result = [self.module_list[col](x[:, col]).unsqueeze(2) for col in range(x.shape[1])]
        return torch.cat(result, dim=2)


class TrainingModuleV2(pl.LightningModule):
    def __init__(self, x, dim_emb, dim_mlp, res_coef=0, dropout_p=0, n_layers=10):
        super().__init__()
        self.num_feature = x.shape[1]
        self.dim_emb = dim_emb
        self.embedding = EmbeddingFactory(x, dim_emb)
        self.backbone = MyNetwork(x.shape[1]*dim_emb, dim_mlp, res_coef, dropout_p, n_layers)
        self.loss = nn.BCELoss()
        self.accuracy = Accuracy()
        
    def forward(self, x):
        x = self.embedding(x)
        #print("x_shape ",x.shape)
        x = x.reshape(-1,self.num_feature*self.dim_emb)
        return self.backbone(x)
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        x = self.embedding(x)
        x = rearrange(x, "b h e -> b (h e)")
        x = self.backbone(x)
        loss = self.loss(x, y.type(torch.float32))
        acc = self.accuracy(x, y)
        self.log("Validation loss", loss)
        self.log("Validation acc", acc)
        return loss, acc
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        x = self.embedding(x)
        x = rearrange(x, "b h e -> b (h e)")
        x = self.backbone(x)
        loss = self.loss(x, y.type(torch.float32))
        acc = self.accuracy(x, y)
        self.log("Training loss", loss)
        self.log("Training acc", acc)
        
        return loss
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer